################################
#
# input: a Newick-formatted tree
#
# method: bins all branch lengths into 0.05-sized bins and counts
# the proportion of BLs in each bin
#
# output: a CRAN script and a PDF plot (generated from the CRAN script)
#
#

import os
import re
import sys

tree_path = sys.argv[1]

fin = open(tree_path, "r")
tree = fin.readlines()[0]
fin.close()

#print "Tree = ", tree

delims = ["\)", "\(", "\:", "\,", "\;"]
for d in delims:
    tree = re.sub(d, "   ", tree)

print tree
tokens = tree.split()

bls = []
for t in tokens:
    x = re.match("\d+\.\d+", t)
    if x != None:
        print t
        bls.append( float(t) )

print bls

# calculate mean
sum = 0.0
for b in bls:
    sum += b
mean = sum / bls.__len__()
mean_str = "%.3f"%mean

# median
bls.sort()
median = bls[len(bls)/2]
median_str = "%.3f"%median

bins = {}
for b in bls:
    rounded_bl = round(b, 2) # round to the nearest 0.01
    if bins.__contains__( rounded_bl ):
        bins[ rounded_bl ] += 1
    else:
        bins[ rounded_bl ] = 1
print bins

#normalize:
max = 0
for b in bins.keys():
    if bins[b] > max:
        max = bins[b]
normalized_bins = {}
for b in bins.keys():
    normalized_bins[b] = float(bins[b]) / max

# print CRAN script
fout = open(tree_path + ".bl_distribution.cran", "w")
lstr = "lengths <- c("
keys = normalized_bins.keys()
keys.sort()
for x in keys:
    this_length = "%.2f"%x
    lstr += this_length + ","
lstr = re.sub(",$", "", lstr)
lstr += ")"
fout.write(lstr + "\n")

pstr = "proportions <- c("
for x in keys:
    this_prop = "%.2f"%normalized_bins[x]
    pstr += this_prop + ","
pstr = re.sub(",$", "", pstr)
pstr += ")"
fout.write(pstr + "\n")

fout.write("pdf(\"" + tree_path + ".bl_distribution.pdf" + "\")\n")
fout.write("plot(lengths, proportions, type='l',log=\"x\",xlab=\"BLs, binned\", ylab=\"proportion\", sub=\"mean BL = " + mean_str.__str__() + ", median = " + median_str.__str__() + "\");\n")
fout.write("dev.off()\n")
fout.close()

os.system("r --no-save < " + tree_path + ".bl_distribution.cran")